from sanic import Sanic, response
from paddleocr import PaddleOCR
import base64
import numpy as np
from sanic.worker.manager import WorkerManager
import logging
import os
from dotenv import load_dotenv

load_dotenv()

use_gpu = os.getenv("OCR_USE_GPU") == "True"

WorkerManager.THRESHOLD = 6000

logger = logging.getLogger('ocr_server')
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logger.addHandler(handler)


logger.info(f"OCR_USE_GPU parameter is set to {use_gpu}")

# 创建 Sanic 应用
app = Sanic("OCRService")


# # load model
# # Paddleocr目前支持中英文、英文、法语、德语、韩语、日语，可以通过修改 lang参数进行切换
# # lang参数依次为`ch`, `en`, `french`, `german`, `korean`, `japan`
# ocr = PaddleOCR(lang="ch",
#                 use_gpu=True,
#                 det_model_dir="/data/lrs/PdfParser/ocr_models/ch_PP-OCRv4_det_infer",
#                 cls_model_dir="/data/lrs/PdfParser/ocr_models/ch_ppocr_mobile_v2.0_cls_infer",
#                 rec_model_dir="/data/lrs/PdfParser/ocr_models/ch_PP-OCRv4_rec_infer")

# 初始化 PaddleOCR 引擎
ocr_engine = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=use_gpu, show_log=False)


# 定义 OCR API 路由
@app.post("/ocr")
async def ocr_request(request):
    # 获取上传的文件
    input = request.json
    img_file = input['img64']
    height = input['height']
    width = input['width']
    channels = input['channels']

    binary_data = base64.b64decode(img_file)
    img_array = np.frombuffer(binary_data, dtype=np.uint8).reshape((height, width, channels))
    logger.info("shape: {}".format(img_array.shape))

    # 无文件上传，返回错误
    if not img_file:
        return response.json({'error': 'No file was uploaded.'}, status=400)

    # 调用 PaddleOCR 进行识别
    res = ocr_engine.ocr(img_array)
    logger.info("ocr result: {}".format(res))

    # 返回识别结果
    return response.json({'results': res})


# """Loader that loads image files."""
# from typing import List, Callable

# from langchain.document_loaders.unstructured import UnstructuredFileLoader
# from unstructured.partition.text import partition_text
# import os
# import fitz
# from tqdm import tqdm
# from typing import Union, Any
# import numpy as np
# import base64


# class UnstructuredPaddlePDFLoader(UnstructuredFileLoader):
#     """Loader that uses unstructured to load image files, such as PNGs and JPGs."""
#     def __init__(
#         self,
#         file_path: Union[str, List[str]],
#         ocr_engine: Callable,
#         mode: str = "single",
#         **unstructured_kwargs: Any,
#     ):
#         """Initialize with file path."""
#         self.ocr_engine = ocr_engine
#         super().__init__(file_path=file_path, mode=mode, **unstructured_kwargs)

#     def _get_elements(self) -> List:
#         def pdf_ocr_txt(filepath, dir_path="tmp_files"):
#             full_dir_path = os.path.join(os.path.dirname(filepath), dir_path)
#             if not os.path.exists(full_dir_path):
#                 os.makedirs(full_dir_path)
#             doc = fitz.open(filepath)
#             txt_file_path = os.path.join(full_dir_path, "{}.txt".format(os.path.split(filepath)[-1]))
#             img_name = os.path.join(full_dir_path, 'tmp.png')
#             with open(txt_file_path, 'w', encoding='utf-8') as fout:
#                 for i in tqdm(range(doc.page_count)):
#                     page = doc.load_page(i)
#                     pix = page.get_pixmap()
#                     img = np.frombuffer(pix.samples, dtype=np.uint8).reshape((pix.h, pix.w, pix.n))

#                     img_data = {"img64": base64.b64encode(img).decode("utf-8"), "height": pix.h, "width": pix.w,
#                                 "channels": pix.n}
#                     result = self.ocr_engine(img_data)
#                     result = [line for line in result if line]
#                     ocr_result = [i[1][0] for line in result for i in line]
#                     fout.write("\n".join(ocr_result))
#             if os.path.exists(img_name):
#                 os.remove(img_name)
#             return txt_file_path

#         txt_file_path = pdf_ocr_txt(self.file_path)
#         return partition_text(filename=txt_file_path, **self.unstructured_kwargs)

# def get_ocr_result(self, image_data: dict):
#     response = requests.post(self.ocr_url, json=image_data)
#     response.raise_for_status()  # 如果请求返回了错误状态码，将会抛出异常
#     return response.json()['results']


# 启动服务
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8010, workers=1, access_log=False)

    # file_path = ''
    # loader = UnstructuredPaddlePDFLoader(file_path, get_ocr_result)
    # docs = loader.load()

    
